import numpy as np
import matplotlib.pyplot as plt

class PointMassModelKalmanFilter:
    def __init__(self, dt, process_noise_std=[0.1, 0.2, 0.1], measurement_noise_std=[0.5, 0.3, 0.2]):
        """
        车辆卡尔曼滤波器初始化
        
        参数:
        dt: 时间步长 (秒)
        process_noise_std: 过程噪声标准差 [位置, 速度, 加速度]
        measurement_noise_std: 观测噪声标准差 [位置, 速度, 加速度]
        """
        self.dt = dt
        
        # 状态向量: [位置y, 速度v, 加速度a]
        self.state_dim = 3
        self.state = np.zeros(self.state_dim)
        # 状态协方差矩阵
        self.P = np.eye(self.state_dim)
        
        # 状态转移矩阵 (基于点质动力学模型)
        # y_{k+1} = y_k + v_k*dt + 0.5*a_k*dt^2
        # v_{k+1} = v_k + a_k*dt
        # a_{k+1} = a_k (假设加速度保持恒定)
        self.F = np.array([
            [1, dt, 0.5*dt**2],
            [0, 1, dt],
            [0, 0, 1]
        ])
        
        # 控制输入矩阵 (加速度作为控制量)
        # 控制量只影响加速度状态
        self.B = np.array([[0], [0], [1]])
        
        # 观测矩阵 (观测所有三个状态)
        self.H = np.eye(self.state_dim)
        
        # 过程噪声协方差矩阵
        self.Q = np.diag(np.array(process_noise_std)**2)
        
        # 观测噪声协方差矩阵
        self.R = np.diag(np.array(measurement_noise_std)**2)
    
    def reset(self, initial_state=None):
        """
        重置滤波器状态
        
        参数:
        initial_state: 初始状态向量 [位置, 速度, 加速度]
        """
        if initial_state is None:
            self.state = np.zeros(self.state_dim)
        else:
            self.state = initial_state.copy()
        
        self.P = np.eye(self.state_dim)
    
    def predict(self, control_acceleration=0):
        """
        预测步骤
        
        参数:
        control_acceleration: 控制加速度输入
        """
        # 状态预测
        control_vector = control_acceleration * self.B.flatten()
        self.state = self.F @ self.state + control_vector
        
        # 协方差预测
        self.P = self.F @ self.P @ self.F.T + self.Q
        
        return self.state.copy()
    
    def update(self, measurement):
        """
        更新步骤
        
        参数:
        measurement: 观测向量 [位置, 速度, 加速度]
        """
        # 计算卡尔曼增益
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        
        # 状态更新
        innovation = measurement - self.H @ self.state
        self.state = self.state + K @ innovation
        
        # 协方差更新
        I = np.eye(self.state_dim)
        self.P = (I - K @ self.H) @ self.P
        
        return self.state.copy()
    
    def filter_step(self, measurement, control_acceleration=0):
        """
        完整的滤波步骤: 预测 + 更新
        
        参数:
        measurement: 观测向量 [位置, 速度, 加速度]
        control_acceleration: 控制加速度输入
        
        返回:
        estimated_state: 估计的状态向量
        """
        self.predict(control_acceleration)
        estimated_state = self.update(measurement)
        return estimated_state
    
    def set_state(self, state):
        """设置状态"""
        self.state = state.copy()
    
    def get_state(self):
        """返回当前状态估计"""
        return self.state.copy()
    
    def get_covariance(self):
        """返回当前协方差矩阵"""
        return self.P.copy()
